import os
import cv2
import numpy as np
import torch
import torch.utils.data
import albumentations as A


class Dataset(torch.utils.data.Dataset):
    def __init__(self, img_ids, img_dir, mask_dir, img_ext, mask_ext, num_classes, transform=None, target_size=None):
        """
        Args:
            img_ids (list): Image ids.
            img_dir: Image file directory.
            mask_dir: Mask file directory.
            img_ext (str): Image file extension.
            mask_ext (str): Mask file extension.
            num_classes (int): Number of classes.
            transform (Compose, optional): Compose transforms of albumentations. Defaults to None.
            target_size (tuple, optional): Target size for resizing images (height, width). Defaults to None.
        """
        self.img_ids = img_ids
        self.img_dir = img_dir
        self.mask_dir = mask_dir
        self.img_ext = img_ext
        self.mask_ext = mask_ext
        self.num_classes = num_classes
        self.transform = transform
        self.target_size = target_size

    def __len__(self):
        return len(self.img_ids)

    def __getitem__(self, idx):
        img_id = self.img_ids[idx]

        # Load image
        img_path = os.path.join(self.img_dir, img_id + self.img_ext)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # Load mask
        mask_path = os.path.join(self.mask_dir, img_id + self.mask_ext)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

        # Resize if target_size is specified
        if self.target_size is not None:
            img = cv2.resize(img, (self.target_size[1], self.target_size[0]), interpolation=cv2.INTER_LINEAR)
            mask = cv2.resize(mask, (self.target_size[1], self.target_size[0]), interpolation=cv2.INTER_NEAREST)

        # Apply transformations
        if self.transform is not None:
            augmented = self.transform(image=img, mask=mask)
            img = augmented['image']
            mask = augmented['mask']

        # Normalize and transpose image
        img = img.astype('float32') / 255
        img = img.transpose(2, 0, 1)

        # Convert mask to tensor and ensure values are within [0, num_classes-1]
        mask = mask.astype('int64')
        mask = torch.from_numpy(mask).long()

        return img, mask, {'img_id': img_id}


class TestDataset(torch.utils.data.Dataset):
    """专门用于测试的数据集类（没有掩码）"""

    def __init__(self, img_ids, img_dir, img_ext='.png', transform=None, target_size=None):
        self.img_ids = img_ids
        self.img_dir = img_dir
        self.img_ext = img_ext
        self.transform = transform
        self.target_size = target_size

    def __len__(self):
        return len(self.img_ids)

    def __getitem__(self, idx):
        img_id = self.img_ids[idx]

        # 加载图像
        img_path = os.path.join(self.img_dir, img_id + self.img_ext)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 调整尺寸
        if self.target_size is not None:
            img = cv2.resize(img, (self.target_size[1], self.target_size[0]),
                             interpolation=cv2.INTER_LINEAR)

        # 应用变换
        if self.transform is not None:
            augmented = self.transform(image=img)
            img = augmented['image']

        # 归一化并转置
        img = img.astype('float32') / 255
        img = img.transpose(2, 0, 1)

        return img, img_id


# 数据验证函数
def validate_dataset(dataset, num_samples=5):
    """
    验证数据集是否正确加载
    Args:
        dataset: 数据集实例
        num_samples: 要检查的样本数量
    """
    print(f"数据集大小: {len(dataset)}")
    print(f"类别数量: {dataset.num_classes}")

    for i in range(min(num_samples, len(dataset))):
        img, mask, meta = dataset[i]
        print(f"样本 {i + 1}:")
        print(f"  图像形状: {img.shape}")
        print(f"  掩码形状: {mask.shape}")
        print(f"  掩码唯一值: {torch.unique(mask)}")
        print(f"  掩码值范围: {mask.min()} - {mask.max()}")
        print(f"  图像ID: {meta['img_id']}")
        print()


# 数据统计函数
def analyze_dataset(dataset):
    """
    分析数据集的类别分布
    Args:
        dataset: 数据集实例
    """
    class_counts = np.zeros(dataset.num_classes)

    for i in range(len(dataset)):
        _, mask, _ = dataset[i]
        for class_idx in range(dataset.num_classes):
            class_counts[class_idx] += (mask == class_idx).sum().item()

    total_pixels = np.sum(class_counts)
    print("数据集类别分布:")
    for class_idx in range(dataset.num_classes):
        percentage = class_counts[class_idx] / total_pixels * 100
        print(f"  类别 {class_idx}: {class_counts[class_idx]} 像素 ({percentage:.2f}%)")

    return class_counts